Skip to content

Use compilation jobs for caching and shared reflection - #311

Merged
maleadt merged 9 commits into
mainfrom
tb/reflection
Sep 9, 2026
Merged

maleadt merged 9 commits into
mainfrom
tb/reflection

Conversation

@maleadt

@maleadt maleadt commented Sep 8, 2026 •

Copy link
Copy Markdown
Member

Changing a kernel's target architecture or tuning hints currently creates separate inference cache entries, even though those settings do not affect Julia inference. Reflection during a launch also reconstructs the compilation request from its function and argument types, losing the launch's options.

Introduce TileJob as the compilation request used by launches, caching, and reflection. Inference is shared across targets and hints, while constant arguments select their own inferred specialization. Compiled binaries are cached per configuration and loaded kernels per CUDA context. Results survive unrelated world-age changes and follow Julia's invalidation when kernel code changes. Compilation and module loading run outside cache locks; simultaneous misses may duplicate work but converge on the published result.

cuTile keeps its compiler and manages its own result storage through CompilerCaching. It shares GPUCompiler's reflection hooks. CUDA.@device_code_typed, CUDA.@device_code_warntype, and CUDA.@device_code_ptx can inspect expressions that launch both CUDA and cuTile kernels, including cached launches. The Julia-level views use the job's optimized inferred code, preserving constant arguments. ct.@device_code_typed now returns a dictionary from job to typed code. cuTile-specific stage macros ignore other backends and retain task-scoped hooks, nesting, and deduplication.

Builds on #302 and requires GPUCompiler 2.6 with the shared reflection helpers. PTX reflection remains unstable because it reads an undocumented CUBIN section.

Validation of the shared reflection integration: 120 checks passed on Julia 1.12.7 with an RTX 5080 across codegen reflection, device compilation, and device core tests. This includes warm mixed-backend launches, constant specializations, nested hooks, concurrent cache misses, and multiple CUDA contexts.

@AntonOresten You convinced me that at least making some part of GPUCompiler.jl reusable isn't too bad.

AntonOresten and others added 7 commits September 6, 2026 16:10
PTX is an implementation detail of `tileiras`, read from an undocumented CUBIN
section; `code_ptx`/`@device_code_ptx` say so and go away with it.

cuTile's `@device_code_sass` duplicated `CUDA.@device_code_sass` with a narrower
scope (only cuTile kernels, by recompiling each signature). The CUDA macro
already shows every loaded module, cuTile's included; `code_sass` remains for
inspecting a signature without a device or under a profiler.

The `@device_code_*` hook now fires once from the launch's compile entry point
rather than from inside the cached emit chain, so the chain checks its own
cached field before recursing and subprogram compilation needs no hook
suppression.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XNHajvrhhKBBUF7w6cDiKr
Share inference across targets and hints, storing compiled binaries per
configuration on the generic or constant-specialized inference result.
Cache loaded kernels per CUDA context and synchronize result publication
without holding a lock during compilation or module loading.

Pass the launch job to task-scoped reflection hooks so inspection preserves
its target, constants, and options, including on cache hits.

Co-authored-by: Tim Besard <tim.besard@gmail.com>
cuTile keeps its own compiler, but reports compilations through
GPUCompiler.compile_hook instead of a hook of its own, and builds its
@device_code_* macros on GPUCompiler.emit_hooked_compilation. TileJob
implements GPUCompiler.code_typed and code_warntype, so GPUCompiler's (and
CUDA.jl's) @device_code_typed and @device_code_warntype observe cuTile
launches too, mixed-backend expressions included.

Requires GPUCompiler 2.6 for the backend-agnostic hook.
With compilations reported through the shared hook, the Julia-level stages
need no cuTile-specific macros: re-export GPUCompiler's, which behave like
their CUDA.jl counterparts (@device_code_typed returns a dictionary from job
to typed code), and keep cuTile's own macros for its own stages.
Display warntype from the inferred TileJob source, including constant arguments. Let CUDA PTX reflection dispatch to cuTile, and filter non-Tile jobs from cuTile-specific stage macros.

Test constant-specialized reflection and warm mixed CUDA/cuTile launches through the shared macros.
Comment thread docs/src/man/debugging.md Outdated
Comment thread src/launch.jl
Comment thread src/compiler/reflection.jl Outdated
Comment thread test/codegen/reflection.jl Outdated
Jobs could be built without a target, deferring the architecture to
codegen and leaving reflection with a different notion of "current
target" than launches. Give every job a target: `sm_arch` defaults to
the current CUDA device, as a launch, and `bytecode_version` to the
preference or the newest version the selected `tileiras` accepts. The
pair is validated when the job is built, so an older bytecode version
needs an explicit `sm_arch` on devices it cannot target, and without a
CUDA device `sm_arch` must be given.

Launches now reject an explicit `sm_arch` other than the device's, since
`tileiras` emits architecture-specific code that only runs there.

The CUDA 13.1 CI job is dropped: 13.1 bytecode cannot target the A100 it
runs on, so no job in the suite can be built there. 13.1 bytecode
emission stays covered by the explicit tests on newer toolchains, which
now pin `sm_arch=v"10.0"` alongside the bytecode version.
`code_typed(f, argtypes)` and `code_structured(f, argtypes)` built a
`TileJob`, whose default bytecode version probes `tileiras`, before
running inference. Without target keywords these stages only need the
inferred code, so run inference directly and leave the job to the
target-dependent stages. Neither a device nor `tileiras` is needed for
them anymore.
@maleadt
maleadt merged commit 2f4e6cc into main Sep 9, 2026
1 check passed
@maleadt
maleadt deleted the tb/reflection branch September 9, 2026 15:35
@AntonOresten

Copy link
Copy Markdown
Collaborator

Thank you very much!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants